from typing import List
from .models import SQLToken, SQLTokenType

class SQLLexer:
    """SQL词法分析器"""

    KEYWORDS = {
        'CREATE', 'TABLE', 'PRIMARY', 'KEY', 'NOT', 'NULL', 'AUTO_INCREMENT',
        'DEFAULT', 'UNIQUE', 'INDEX', 'COMMENT', 'ENGINE', 'CHARSET'
    }

    OPERATORS = {'+', '-', '*', '/', '%', '=', '<', '>', '<=', '>=', '<>', '!='}

    def __init__(self, sql: str):
        """初始化SQL词法分析器

        Args:
            sql: SQL语句
        """
        self.sql = sql
        self.pos = 0
        self.line = 1
        self.column = 1
        self.tokens = []

    def tokenize(self) -> List[SQLToken]:
        """执行词法分析,生成token列表"""
        while self.pos < len(self.sql):
            char = self.sql[self.pos]

            if char.isspace():
                self._handle_whitespace()
            elif char.isalpha() or char == '_':
                self._handle_identifier()
            elif char in '\'"`':
                self._handle_string()
            elif char in '(),;':
                self.tokens.append(SQLToken(SQLTokenType.PUNCTUATION, char, self.line, self.column))
                self.pos += 1
                self.column += 1
            elif char == '-' and self._peek() == '-':
                self._handle_comment()
            elif char == '/' and self._peek() == '*':
                self._handle_multiline_comment()
            elif char.isdigit():
                self._handle_number()
            elif char in self.OPERATORS:
                self._handle_operator()
            else:
                self.pos += 1
                self.column += 1

        return [t for t in self.tokens if t.type != SQLTokenType.WHITESPACE]

    def _handle_operator(self):
        start = self.pos
        start_column = self.column
        current_op = self.sql[self.pos]

        # 检查是否是双字符运算符
        if self.pos + 1 < len(self.sql):
            next_char = self.sql[self.pos + 1]
            double_op = current_op + next_char
            if double_op in self.OPERATORS:
                current_op = double_op
                self.pos += 2
                self.column += 2
            else:
                self.pos += 1
                self.column += 1
        else:
            self.pos += 1
            self.column += 1

        self.tokens.append(SQLToken(SQLTokenType.OPERATOR, current_op, self.line, start_column))

    def _peek(self, offset: int = 1) -> str:
        if self.pos + offset < len(self.sql):
            return self.sql[self.pos + offset]
        return ''

    def _handle_whitespace(self):
        start = self.pos
        start_column = self.column

        while self.pos < len(self.sql) and self.sql[self.pos].isspace():
            if self.sql[self.pos] == '\n':
                self.line += 1
                self.column = 1
            else:
                self.column += 1
            self.pos += 1

        self.tokens.append(SQLToken(SQLTokenType.WHITESPACE, self.sql[start:self.pos], self.line, start_column))

    def _handle_identifier(self):
        start = self.pos
        start_column = self.column

        while self.pos < len(self.sql) and (self.sql[self.pos].isalnum() or self.sql[self.pos] == '_'):
            self.pos += 1
            self.column += 1

        value = self.sql[start:self.pos]
        token_type = SQLTokenType.KEYWORD if value.upper() in self.KEYWORDS else SQLTokenType.IDENTIFIER
        self.tokens.append(SQLToken(token_type, value, self.line, start_column))

    def _handle_string(self):
        quote = self.sql[self.pos]
        start = self.pos
        start_column = self.column
        self.pos += 1
        self.column += 1

        while self.pos < len(self.sql) and self.sql[self.pos] != quote:
            if self.sql[self.pos] == '\n':
                self.line += 1
                self.column = 1
            else:
                self.column += 1
            self.pos += 1

        if self.pos < len(self.sql):  # 包含结束引号
            self.pos += 1
            self.column += 1

        self.tokens.append(SQLToken(SQLTokenType.STRING, self.sql[start:self.pos], self.line, start_column))

    def _handle_punctuation(self):
        self.tokens.append(SQLToken(SQLTokenType.PUNCTUATION, self.sql[self.pos], self.line, self.column))
        self.pos += 1
        self.column += 1

    def _handle_comment(self):
        start = self.pos
        start_column = self.column

        while self.pos < len(self.sql) and self.sql[self.pos] != '\n':
            self.pos += 1
            self.column += 1

        self.tokens.append(SQLToken(SQLTokenType.COMMENT, self.sql[start:self.pos], self.line, start_column))

    def _handle_multiline_comment(self):
        start = self.pos
        start_column = self.column
        self.pos += 2  # 跳过 /*
        self.column += 2

        while self.pos < len(self.sql) - 1 and not (self.sql[self.pos] == '*' and self.sql[self.pos + 1] == '/'):
            if self.sql[self.pos] == '\n':
                self.line += 1
                self.column = 1
            else:
                self.column += 1
            self.pos += 1

        if self.pos < len(self.sql) - 1:  # 包含结束标记
            self.pos += 2
            self.column += 2

        self.tokens.append(SQLToken(SQLTokenType.COMMENT, self.sql[start:self.pos], self.line, start_column))

    def _handle_number(self):
        start = self.pos
        start_column = self.column

        while self.pos < len(self.sql) and (self.sql[self.pos].isdigit() or self.sql[self.pos] == '.'):
            self.pos += 1
            self.column += 1

        self.tokens.append(SQLToken(SQLTokenType.NUMBER, self.sql[start:self.pos], self.line, start_column))